package cn.veliger.springai.controller;

import cn.veliger.springai.advisor.CusCallAroundAdvisor;
import cn.veliger.springai.advisor.CusRetrievalAugmentationAdvisor;
import cn.veliger.springai.service.DocumentService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.evaluation.RelevancyEvaluator;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.evaluation.EvaluationRequest;
import org.springframework.ai.evaluation.EvaluationResponse;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.Map;

/**
 * @author veliger
 * create on 2025-06-28
 **/
@RestController
@RequestMapping
@RequiredArgsConstructor
public class RagController {
    private final EmbeddingModel embeddingModel;
    private final DocumentService documentService;
    private final ChatClient chatClient;
    private final VectorStore vectorStore;
    private final ChatModel chatModel;
    private final ChatClient memoryChatClient;

    /**
     * 向量化
     *
     * @param msg
     * @return
     */
    @GetMapping("/embedding")
    public float[] embedding(@RequestParam String msg) {
        return embeddingModel.embed(msg);
    }

    /**
     * 文档分割
     *
     * @return
     */
    @GetMapping("/documentEmbedding")
    public List<Document> documentEmbedding() {
        return documentService.loadText();
    }

    /**
     * 向量查询
     *
     * @param msg
     * @return
     */
    @GetMapping("/documentSearch")
    public List<Document> documentSearch(@RequestParam String msg) {
        long currentTimeMillis = System.currentTimeMillis();
        List<Document> search = documentService.search(msg);
        System.out.println("耗时：" + (System.currentTimeMillis() - currentTimeMillis));
        return search;
    }

    /**
     * rag聊天
     *
     * @param msg
     * @return
     */
    @GetMapping("/ragChat")
    public String ragChat(@RequestParam String msg) {

        // 向量搜索
        List<Document> documentList = documentService.search(msg);

        // 提示词模板
        PromptTemplate promptTemplate = new PromptTemplate("{userMessage}\n\n 用以下信息回答问题:\n {contents}");

        // 组装提示词
        Prompt prompt = promptTemplate.create(Map.of("userMessage", msg, "contents", documentList));

        // 调用大模型
        return chatClient.prompt(prompt)
                .advisors(new CusCallAroundAdvisor())
                .call().content();
    }

    /**
     * QueryTransformer
     *  - CompressionQueryTransformer 对提问进行压缩
     *  - RewriteQueryTransformer 对用户提问进行重写
     *  - TranslationQueryTransformer 对用户提问进行翻译
     * QueryExpander
     *  - MultiQueryExpander 对用户提问进行扩展
     *  DocumentRetriever
     *  - VectorStoreDocumentRetriever 从向量数据库中进行文档检索
     *  DocumentJoiner
     *  - ConcatenationDocumentJoiner 对多个Document进行合并，会根据Document的id进行去重
     *  QueryAugmenter
     *  - ContextualQueryAugmenter 增强查询，基于用户问题和检索结果进行增强。
     */


    /**
     * RetrievalAugmentationAdvisor底层流程
     * 1. 先基于用户当前问题和会话记录，构造一个Query对象
     * 2. 利用QueryTransformer对Query进行转换
     * 3. 利用QueryExpander对Query进行扩展，扩展出多个Query
     * 4. 并行根据多个Query利用DocumentRetriever进行向量数据库检索
     * 5. 把所有检索结果进行去重合并
     * 6. 利用QueryAugmenter对原始Query以及检索结果进行增强，得到最终的Query
     * 7. 基于最终的Query构造并返回新的AdvisedRequest对象
     */
    /**
     * @param msg
     * @return
     */
    @GetMapping("/qa")
    public String qa(@RequestParam("chatId") String chatId, @RequestParam String msg) {
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.50)
                        .vectorStore(vectorStore)
                        .build())
                .build();
        CusRetrievalAugmentationAdvisor cusRetrievalAugmentationAdvisor = CusRetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.50)
                        .topK(1)
                        .vectorStore(vectorStore)
                        .build())
                .build();
      /*  QuestionAnswerAdvisor questionAnswerAdvisor = new QuestionAnswerAdvisor(vectorStore);

        var qaAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(SearchRequest.builder().similarityThreshold(0.5d).topK(1).build())
                .build();*/
        /*Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.50)
                        .topK(5)
                        .vectorStore(vectorStore)
                        .build())
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .allowEmptyContext(true)
                        .build())
                .build();*/
        return memoryChatClient
                .prompt()
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .advisors(retrievalAugmentationAdvisor, new CusCallAroundAdvisor())
                .user(msg)
                .call()
                .content();
    }

    /**
     * 幻觉评估
     *
     * @param message
     * @return
     */
    @GetMapping("/evaluation")
    public EvaluationResponse evaluation(@RequestParam String message) {
        // 向量搜索
        List<Document> documentList = documentService.search(message);
        // 提示词模板
        PromptTemplate promptTemplate = new PromptTemplate("{userMessage}\n\n 用以下信息回答问题:\n {contents}");
        // 组装提示词
        Prompt prompt = promptTemplate.create(Map.of("userMessage", message, "contents", documentList));
        // 调用大模型
        String result = chatClient.prompt(prompt).call().content();
        //        String result = "我是机器人";
        var relevancyEvaluator = new RelevancyEvaluator(ChatClient.builder(chatModel));
        // 评估是否产生了幻觉
        EvaluationRequest evaluationRequest = new EvaluationRequest(message,
                documentList, result);
        EvaluationResponse evaluationResponse = relevancyEvaluator.evaluate(evaluationRequest);

        return evaluationResponse;
    }

}
